{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# Model interpretability with Integrated Gradients\n",
    "\n",
    "**Author:** [A_K_Nain](https://twitter.com/A_K_Nain)<br>\n",
    "**Date created:** 2020/06/02<br>\n",
    "**Last modified:** 2020/06/02<br>\n",
    "**Description:** How to obtain integrated gradients for a classification model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Integrated Gradients\n",
    "\n",
    "[Integrated Gradients](https://arxiv.org/abs/1703.01365) is a technique for\n",
    "attributing a classification model's prediction to its input features. It is\n",
    "a model interpretability technique: you can use it to visualize the relationship\n",
    "between input features and model predictions.\n",
    "\n",
    "Integrated Gradients is a variation on computing\n",
    "the gradient of the prediction output with regard to features of the input.\n",
    "To compute integrated gradients, we need to perform the following steps:\n",
    "\n",
    "1. Identify the input and the output. In our case, the input is an image and the\n",
    "output is the last layer of our model (dense layer with softmax activation).\n",
    "\n",
    "2. Compute which features are important to a neural network\n",
    "when making a prediction on a particular data point. To identify these features, we\n",
    "need to choose a baseline input. A baseline input can be a black image (all pixel\n",
    "values set to zero) or random noise. The shape of the baseline input needs to be\n",
    "the same as our input image, e.g. (299, 299, 3).\n",
    "\n",
    "3. Interpolate the baseline for a given number of steps. The number of steps represents\n",
    "the steps we need in the gradient approximation for a given input image. The number of\n",
    "steps is a hyperparameter. The authors recommend using anywhere between\n",
    "20 and 1000 steps.\n",
    "\n",
    "4. Preprocess these interpolated images and do a forward pass.\n",
    "5. Get the gradients for these interpolated images.\n",
    "6. Approximate the gradients integral using the trapezoidal rule.\n",
    "\n",
    "To read in-depth about integrated gradients and why this method works,\n",
    "consider reading this excellent\n",
    "[article](https://distill.pub/2020/attribution-baselines/).\n",
    "\n",
    "**References:**\n",
    "\n",
    "- Integrated Gradients original [paper](https://arxiv.org/abs/1703.01365)\n",
    "- [Original implementation](https://github.com/ankurtaly/Integrated-Gradients)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Setup\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import ndimage\n",
    "from IPython.display import Image\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.applications import xception\n",
    "\n",
    "# Size of the input image\n",
    "img_size = (299, 299, 3)\n",
    "\n",
    "# Load Xception model with imagenet weights\n",
    "model = xception.Xception(weights=\"imagenet\")\n",
    "\n",
    "# The local path to our target image\n",
    "img_path = keras.utils.get_file(\"elephant.jpg\", \"https://i.imgur.com/Bvro0YD.png\")\n",
    "display(Image(img_path))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Integrated Gradients algorithm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "\n",
    "def get_img_array(img_path, size=(299, 299)):\n",
    "    # `img` is a PIL image of size 299x299\n",
    "    img = keras.preprocessing.image.load_img(img_path, target_size=size)\n",
    "    # `array` is a float32 Numpy array of shape (299, 299, 3)\n",
    "    array = keras.preprocessing.image.img_to_array(img)\n",
    "    # We add a dimension to transform our array into a \"batch\"\n",
    "    # of size (1, 299, 299, 3)\n",
    "    array = np.expand_dims(array, axis=0)\n",
    "    return array\n",
    "\n",
    "\n",
    "def get_gradients(img_input, top_pred_idx):\n",
    "    \"\"\"Computes the gradients of outputs w.r.t input image.\n",
    "\n",
    "    Args:\n",
    "        img_input: 4D image tensor\n",
    "        top_pred_idx: Predicted label for the input image\n",
    "\n",
    "    Returns:\n",
    "        Gradients of the predictions w.r.t img_input\n",
    "    \"\"\"\n",
    "    images = tf.cast(img_input, tf.float32)\n",
    "\n",
    "    with tf.GradientTape() as tape:\n",
    "        tape.watch(images)\n",
    "        preds = model(images)\n",
    "        top_class = preds[:, top_pred_idx]\n",
    "\n",
    "    grads = tape.gradient(top_class, images)\n",
    "    return grads\n",
    "\n",
    "\n",
    "def get_integrated_gradients(img_input, top_pred_idx, baseline=None, num_steps=50):\n",
    "    \"\"\"Computes Integrated Gradients for a predicted label.\n",
    "\n",
    "    Args:\n",
    "        img_input (ndarray): Original image\n",
    "        top_pred_idx: Predicted label for the input image\n",
    "        baseline (ndarray): The baseline image to start with for interpolation\n",
    "        num_steps: Number of interpolation steps between the baseline\n",
    "            and the input used in the computation of integrated gradients. These\n",
    "            steps along determine the integral approximation error. By default,\n",
    "            num_steps is set to 50.\n",
    "\n",
    "    Returns:\n",
    "        Integrated gradients w.r.t input image\n",
    "    \"\"\"\n",
    "    # If baseline is not provided, start with a black image\n",
    "    # having same size as the input image.\n",
    "    if baseline is None:\n",
    "        baseline = np.zeros(img_size).astype(np.float32)\n",
    "    else:\n",
    "        baseline = baseline.astype(np.float32)\n",
    "\n",
    "    # 1. Do interpolation.\n",
    "    img_input = img_input.astype(np.float32)\n",
    "    interpolated_image = [\n",
    "        baseline + (step / num_steps) * (img_input - baseline)\n",
    "        for step in range(num_steps + 1)\n",
    "    ]\n",
    "    interpolated_image = np.array(interpolated_image).astype(np.float32)\n",
    "\n",
    "    # 2. Preprocess the interpolated images\n",
    "    interpolated_image = xception.preprocess_input(interpolated_image)\n",
    "\n",
    "    # 3. Get the gradients\n",
    "    grads = []\n",
    "    for i, img in enumerate(interpolated_image):\n",
    "        img = tf.expand_dims(img, axis=0)\n",
    "        grad = get_gradients(img, top_pred_idx=top_pred_idx)\n",
    "        grads.append(grad[0])\n",
    "    grads = tf.convert_to_tensor(grads, dtype=tf.float32)\n",
    "\n",
    "    # 4. Approximate the integral using the trapezoidal rule\n",
    "    grads = (grads[:-1] + grads[1:]) / 2.0\n",
    "    avg_grads = tf.reduce_mean(grads, axis=0)\n",
    "\n",
    "    # 5. Calculate integrated gradients and return\n",
    "    integrated_grads = (img_input - baseline) * avg_grads\n",
    "    return integrated_grads\n",
    "\n",
    "\n",
    "def random_baseline_integrated_gradients(\n",
    "    img_input, top_pred_idx, num_steps=50, num_runs=2\n",
    "):\n",
    "    \"\"\"Generates a number of random baseline images.\n",
    "\n",
    "    Args:\n",
    "        img_input (ndarray): 3D image\n",
    "        top_pred_idx: Predicted label for the input image\n",
    "        num_steps: Number of interpolation steps between the baseline\n",
    "            and the input used in the computation of integrated gradients. These\n",
    "            steps along determine the integral approximation error. By default,\n",
    "            num_steps is set to 50.\n",
    "        num_runs: number of baseline images to generate\n",
    "\n",
    "    Returns:\n",
    "        Averaged integrated gradients for `num_runs` baseline images\n",
    "    \"\"\"\n",
    "    # 1. List to keep track of Integrated Gradients (IG) for all the images\n",
    "    integrated_grads = []\n",
    "\n",
    "    # 2. Get the integrated gradients for all the baselines\n",
    "    for run in range(num_runs):\n",
    "        baseline = np.random.random(img_size) * 255\n",
    "        igrads = get_integrated_gradients(\n",
    "            img_input=img_input,\n",
    "            top_pred_idx=top_pred_idx,\n",
    "            baseline=baseline,\n",
    "            num_steps=num_steps,\n",
    "        )\n",
    "        integrated_grads.append(igrads)\n",
    "\n",
    "    # 3. Return the average integrated gradients for the image\n",
    "    integrated_grads = tf.convert_to_tensor(integrated_grads)\n",
    "    return tf.reduce_mean(integrated_grads, axis=0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Helper class for visualizing gradients and integrated gradients\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "\n",
    "class GradVisualizer:\n",
    "    \"\"\"Plot gradients of the outputs w.r.t an input image.\"\"\"\n",
    "\n",
    "    def __init__(self, positive_channel=None, negative_channel=None):\n",
    "        if positive_channel is None:\n",
    "            self.positive_channel = [0, 255, 0]\n",
    "        else:\n",
    "            self.positive_channel = positive_channel\n",
    "\n",
    "        if negative_channel is None:\n",
    "            self.negative_channel = [255, 0, 0]\n",
    "        else:\n",
    "            self.negative_channel = negative_channel\n",
    "\n",
    "    def apply_polarity(self, attributions, polarity):\n",
    "        if polarity == \"positive\":\n",
    "            return np.clip(attributions, 0, 1)\n",
    "        else:\n",
    "            return np.clip(attributions, -1, 0)\n",
    "\n",
    "    def apply_linear_transformation(\n",
    "        self,\n",
    "        attributions,\n",
    "        clip_above_percentile=99.9,\n",
    "        clip_below_percentile=70.0,\n",
    "        lower_end=0.2,\n",
    "    ):\n",
    "        # 1. Get the thresholds\n",
    "        m = self.get_thresholded_attributions(\n",
    "            attributions, percentage=100 - clip_above_percentile\n",
    "        )\n",
    "        e = self.get_thresholded_attributions(\n",
    "            attributions, percentage=100 - clip_below_percentile\n",
    "        )\n",
    "\n",
    "        # 2. Transform the attributions by a linear function f(x) = a*x + b such that\n",
    "        # f(m) = 1.0 and f(e) = lower_end\n",
    "        transformed_attributions = (1 - lower_end) * (np.abs(attributions) - e) / (\n",
    "            m - e\n",
    "        ) + lower_end\n",
    "\n",
    "        # 3. Make sure that the sign of transformed attributions is the same as original attributions\n",
    "        transformed_attributions *= np.sign(attributions)\n",
    "\n",
    "        # 4. Only keep values that are bigger than the lower_end\n",
    "        transformed_attributions *= transformed_attributions >= lower_end\n",
    "\n",
    "        # 5. Clip values and return\n",
    "        transformed_attributions = np.clip(transformed_attributions, 0.0, 1.0)\n",
    "        return transformed_attributions\n",
    "\n",
    "    def get_thresholded_attributions(self, attributions, percentage):\n",
    "        if percentage == 100.0:\n",
    "            return np.min(attributions)\n",
    "\n",
    "        # 1. Flatten the attributions\n",
    "        flatten_attr = attributions.flatten()\n",
    "\n",
    "        # 2. Get the sum of the attributions\n",
    "        total = np.sum(flatten_attr)\n",
    "\n",
    "        # 3. Sort the attributions from largest to smallest.\n",
    "        sorted_attributions = np.sort(np.abs(flatten_attr))[::-1]\n",
    "\n",
    "        # 4. Calculate the percentage of the total sum that each attribution\n",
    "        # and the values about it contribute.\n",
    "        cum_sum = 100.0 * np.cumsum(sorted_attributions) / total\n",
    "\n",
    "        # 5. Threshold the attributions by the percentage\n",
    "        indices_to_consider = np.where(cum_sum >= percentage)[0][0]\n",
    "\n",
    "        # 6. Select the desired attributions and return\n",
    "        attributions = sorted_attributions[indices_to_consider]\n",
    "        return attributions\n",
    "\n",
    "    def binarize(self, attributions, threshold=0.001):\n",
    "        return attributions > threshold\n",
    "\n",
    "    def morphological_cleanup_fn(self, attributions, structure=np.ones((4, 4))):\n",
    "        closed = ndimage.grey_closing(attributions, structure=structure)\n",
    "        opened = ndimage.grey_opening(closed, structure=structure)\n",
    "        return opened\n",
    "\n",
    "    def draw_outlines(\n",
    "        self, attributions, percentage=90, connected_component_structure=np.ones((3, 3))\n",
    "    ):\n",
    "        # 1. Binarize the attributions.\n",
    "        attributions = self.binarize(attributions)\n",
    "\n",
    "        # 2. Fill the gaps\n",
    "        attributions = ndimage.binary_fill_holes(attributions)\n",
    "\n",
    "        # 3. Compute connected components\n",
    "        connected_components, num_comp = ndimage.measurements.label(\n",
    "            attributions, structure=connected_component_structure\n",
    "        )\n",
    "\n",
    "        # 4. Sum up the attributions for each component\n",
    "        total = np.sum(attributions[connected_components > 0])\n",
    "        component_sums = []\n",
    "        for comp in range(1, num_comp + 1):\n",
    "            mask = connected_components == comp\n",
    "            component_sum = np.sum(attributions[mask])\n",
    "            component_sums.append((component_sum, mask))\n",
    "\n",
    "        # 5. Compute the percentage of top components to keep\n",
    "        sorted_sums_and_masks = sorted(component_sums, key=lambda x: x[0], reverse=True)\n",
    "        sorted_sums = list(zip(*sorted_sums_and_masks))[0]\n",
    "        cumulative_sorted_sums = np.cumsum(sorted_sums)\n",
    "        cutoff_threshold = percentage * total / 100\n",
    "        cutoff_idx = np.where(cumulative_sorted_sums >= cutoff_threshold)[0][0]\n",
    "        if cutoff_idx > 2:\n",
    "            cutoff_idx = 2\n",
    "\n",
    "        # 6. Set the values for the kept components\n",
    "        border_mask = np.zeros_like(attributions)\n",
    "        for i in range(cutoff_idx + 1):\n",
    "            border_mask[sorted_sums_and_masks[i][1]] = 1\n",
    "\n",
    "        # 7. Make the mask hollow and show only the border\n",
    "        eroded_mask = ndimage.binary_erosion(border_mask, iterations=1)\n",
    "        border_mask[eroded_mask] = 0\n",
    "\n",
    "        # 8. Return the outlined mask\n",
    "        return border_mask\n",
    "\n",
    "    def process_grads(\n",
    "        self,\n",
    "        image,\n",
    "        attributions,\n",
    "        polarity=\"positive\",\n",
    "        clip_above_percentile=99.9,\n",
    "        clip_below_percentile=0,\n",
    "        morphological_cleanup=False,\n",
    "        structure=np.ones((3, 3)),\n",
    "        outlines=False,\n",
    "        outlines_component_percentage=90,\n",
    "        overlay=True,\n",
    "    ):\n",
    "        if polarity not in [\"positive\", \"negative\"]:\n",
    "            raise ValueError(\n",
    "                f\"\"\" Allowed polarity values: 'positive' or 'negative'\n",
    "                                    but provided {polarity}\"\"\"\n",
    "            )\n",
    "        if clip_above_percentile < 0 or clip_above_percentile > 100:\n",
    "            raise ValueError(\"clip_above_percentile must be in [0, 100]\")\n",
    "\n",
    "        if clip_below_percentile < 0 or clip_below_percentile > 100:\n",
    "            raise ValueError(\"clip_below_percentile must be in [0, 100]\")\n",
    "\n",
    "        # 1. Apply polarity\n",
    "        if polarity == \"positive\":\n",
    "            attributions = self.apply_polarity(attributions, polarity=polarity)\n",
    "            channel = self.positive_channel\n",
    "        else:\n",
    "            attributions = self.apply_polarity(attributions, polarity=polarity)\n",
    "            attributions = np.abs(attributions)\n",
    "            channel = self.negative_channel\n",
    "\n",
    "        # 2. Take average over the channels\n",
    "        attributions = np.average(attributions, axis=2)\n",
    "\n",
    "        # 3. Apply linear transformation to the attributions\n",
    "        attributions = self.apply_linear_transformation(\n",
    "            attributions,\n",
    "            clip_above_percentile=clip_above_percentile,\n",
    "            clip_below_percentile=clip_below_percentile,\n",
    "            lower_end=0.0,\n",
    "        )\n",
    "\n",
    "        # 4. Cleanup\n",
    "        if morphological_cleanup:\n",
    "            attributions = self.morphological_cleanup_fn(\n",
    "                attributions, structure=structure\n",
    "            )\n",
    "        # 5. Draw the outlines\n",
    "        if outlines:\n",
    "            attributions = self.draw_outlines(\n",
    "                attributions, percentage=outlines_component_percentage\n",
    "            )\n",
    "\n",
    "        # 6. Expand the channel axis and convert to RGB\n",
    "        attributions = np.expand_dims(attributions, 2) * channel\n",
    "\n",
    "        # 7.Superimpose on the original image\n",
    "        if overlay:\n",
    "            attributions = np.clip((attributions * 0.8 + image), 0, 255)\n",
    "        return attributions\n",
    "\n",
    "    def visualize(\n",
    "        self,\n",
    "        image,\n",
    "        gradients,\n",
    "        integrated_gradients,\n",
    "        polarity=\"positive\",\n",
    "        clip_above_percentile=99.9,\n",
    "        clip_below_percentile=0,\n",
    "        morphological_cleanup=False,\n",
    "        structure=np.ones((3, 3)),\n",
    "        outlines=False,\n",
    "        outlines_component_percentage=90,\n",
    "        overlay=True,\n",
    "        figsize=(15, 8),\n",
    "    ):\n",
    "        # 1. Make two copies of the original image\n",
    "        img1 = np.copy(image)\n",
    "        img2 = np.copy(image)\n",
    "\n",
    "        # 2. Process the normal gradients\n",
    "        grads_attr = self.process_grads(\n",
    "            image=img1,\n",
    "            attributions=gradients,\n",
    "            polarity=polarity,\n",
    "            clip_above_percentile=clip_above_percentile,\n",
    "            clip_below_percentile=clip_below_percentile,\n",
    "            morphological_cleanup=morphological_cleanup,\n",
    "            structure=structure,\n",
    "            outlines=outlines,\n",
    "            outlines_component_percentage=outlines_component_percentage,\n",
    "            overlay=overlay,\n",
    "        )\n",
    "\n",
    "        # 3. Process the integrated gradients\n",
    "        igrads_attr = self.process_grads(\n",
    "            image=img2,\n",
    "            attributions=integrated_gradients,\n",
    "            polarity=polarity,\n",
    "            clip_above_percentile=clip_above_percentile,\n",
    "            clip_below_percentile=clip_below_percentile,\n",
    "            morphological_cleanup=morphological_cleanup,\n",
    "            structure=structure,\n",
    "            outlines=outlines,\n",
    "            outlines_component_percentage=outlines_component_percentage,\n",
    "            overlay=overlay,\n",
    "        )\n",
    "\n",
    "        _, ax = plt.subplots(1, 3, figsize=figsize)\n",
    "        ax[0].imshow(image)\n",
    "        ax[1].imshow(grads_attr.astype(np.uint8))\n",
    "        ax[2].imshow(igrads_attr.astype(np.uint8))\n",
    "\n",
    "        ax[0].set_title(\"Input\")\n",
    "        ax[1].set_title(\"Normal gradients\")\n",
    "        ax[2].set_title(\"Integrated gradients\")\n",
    "        plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Let's test-drive it\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "# 1. Convert the image to numpy array\n",
    "img = get_img_array(img_path)\n",
    "\n",
    "# 2. Keep a copy of the original image\n",
    "orig_img = np.copy(img[0]).astype(np.uint8)\n",
    "\n",
    "# 3. Preprocess the image\n",
    "img_processed = tf.cast(xception.preprocess_input(img), dtype=tf.float32)\n",
    "\n",
    "# 4. Get model predictions\n",
    "preds = model.predict(img_processed)\n",
    "top_pred_idx = tf.argmax(preds[0])\n",
    "print(\"Predicted:\", top_pred_idx, xception.decode_predictions(preds, top=1)[0])\n",
    "\n",
    "# 5. Get the gradients of the last layer for the predicted label\n",
    "grads = get_gradients(img_processed, top_pred_idx=top_pred_idx)\n",
    "\n",
    "# 6. Get the integrated gradients\n",
    "igrads = random_baseline_integrated_gradients(\n",
    "    np.copy(orig_img), top_pred_idx=top_pred_idx, num_steps=50, num_runs=2\n",
    ")\n",
    "\n",
    "# 7. Process the gradients and plot\n",
    "vis = GradVisualizer()\n",
    "vis.visualize(\n",
    "    image=orig_img,\n",
    "    gradients=grads[0].numpy(),\n",
    "    integrated_gradients=igrads.numpy(),\n",
    "    clip_above_percentile=99,\n",
    "    clip_below_percentile=0,\n",
    ")\n",
    "\n",
    "vis.visualize(\n",
    "    image=orig_img,\n",
    "    gradients=grads[0].numpy(),\n",
    "    integrated_gradients=igrads.numpy(),\n",
    "    clip_above_percentile=95,\n",
    "    clip_below_percentile=28,\n",
    "    morphological_cleanup=True,\n",
    "    outlines=True,\n",
    ")\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "integrated_gradients",
   "private_outputs": false,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}